# -*- coding: utf-8 -*-
"""
项目名称：cotton_test
文件名称：models.py
创建时间：2025-06-08 20:47:50

系统用户：wyl
作　　者：無以菱
联系邮箱：huangjing2001.guet@qq.com
功能描述：棉花叶片疾病分类模型定义
本模块定义了两种深度学习模型用于棉花叶片疾病分类：
- 1. SimpleMLP: 简单的多层感知机(BP神经网络)
- 2. CustomCNN: 自定义卷积神经网络
这些模型用于区分健康的棉花叶片和患病的棉花叶片。
"""


import torch
import torch.nn as nn

class SimpleMLP(nn.Module):
    """
    简单的多层感知机(BP神经网络)
    
    这是一个基础的全连接神经网络，包含一个隐藏层。
    网络结构：输入层 -> 隐藏层(128个神经元) -> 输出层
    """
    def __init__(self, num_classes=2):
        """
        初始化SimpleMLP模型
        
        参数:
            num_classes (int): 分类类别数，默认为2（健康和患病棉花叶片）
        """
        super(SimpleMLP, self).__init__()
        # 计算输入特征维度：3通道 * 224高 * 224宽
        self.input_size = 3 * 224 * 224  # 输入图像大小
        self.hidden_size = 128  # 隐藏层大小
        self.num_classes = num_classes  # 类别数
        
        # 定义网络层
        self.fc1 = nn.Linear(self.input_size, self.hidden_size)  # 第一个全连接层：输入 -> 隐藏层
        self.relu = nn.ReLU()  # ReLU激活函数：引入非线性
        self.fc2 = nn.Linear(self.hidden_size, self.num_classes)  # 输出层：隐藏层 -> 类别预测

    def forward(self, x):
        """
        前向传播
        
        参数:
            x (torch.Tensor): 输入张量
            
        返回:
            torch.Tensor: 模型输出
        """
        x = x.view(-1, self.input_size)  # 将输入展平
        x = self.relu(self.fc1(x))  # 第一层全连接+激活
        x = self.fc2(x)  # 输出层
        return x


class CustomCNN(nn.Module):
    """
    自定义卷积神经网络
    
    这是一个深度卷积神经网络，用于图像分类任务。
    网络结构包括多个卷积层、池化层和全连接层，设计用于提取图像的层次特征。
    
    网络架构：
    1. 卷积块1: Conv(3->64) -> ReLU -> Conv(64->128) -> ReLU -> MaxPool
    2. 卷积块2: Conv(128->256) -> ReLU -> Conv(256->256) -> ReLU -> MaxPool
    3. 卷积块3: Conv(256->512) -> ReLU -> MaxPool
    4. 全连接层: FC(512*28*28->4096) -> ReLU -> Dropout -> FC(4096->1000) -> ReLU -> Dropout -> FC(1000->num_classes)
    """
    def __init__(self, num_classes=2):
        """
        初始化CustomCNN模型
        
        参数:
            num_classes (int): 分类类别数，默认为2（健康和患病棉花叶片）
        """
        super(CustomCNN, self).__init__()
        
        # 第一个卷积块 - 卷积层1 + ReLU
        # 输入: 3通道图像 -> 输出: 64通道特征图
        # kernel_size=3: 3x3卷积核, padding=1: 保持特征图尺寸不变
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.relu1 = nn.ReLU(inplace=True)  # inplace=True可以节省内存
        
        # 第一个卷积块 - 卷积层2 + ReLU
        # 输入: 64通道特征图 -> 输出: 128通道特征图
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.relu2 = nn.ReLU(inplace=True)
        
        # 第一个池化层 - 最大池化
        # 将特征图尺寸减半: 224x224 -> 112x112
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  # kernel_size=2: 2x2窗口, stride=2: 步长为2
        
        # 第二个卷积块 - 卷积层3 + ReLU
        # 输入: 128通道特征图 -> 输出: 256通道特征图
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.relu3 = nn.ReLU(inplace=True)
        
        # 第二个卷积块 - 卷积层4 + ReLU
        # 输入: 256通道特征图 -> 输出: 256通道特征图 (保持通道数不变)
        self.conv4 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.relu4 = nn.ReLU(inplace=True)
        
        # 第二个池化层 - 最大池化
        # 将特征图尺寸减半: 112x112 -> 56x56
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # 第三个卷积块 - 卷积层5 + ReLU
        # 输入: 256通道特征图 -> 输出: 512通道特征图
        self.conv5 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.relu5 = nn.ReLU(inplace=True)
        
        # 第三个池化层 - 最大池化
        # 将特征图尺寸减半: 56x56 -> 28x28
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # 特征提取部分 - 将所有卷积和池化层组合成一个Sequential模块
        # 这样可以更方便地进行前向传播计算
        self.features = nn.Sequential(
            self.conv1,  # 第一个卷积层
            self.relu1,  # 第一个ReLU激活
            self.conv2,  # 第二个卷积层
            self.relu2,  # 第二个ReLU激活
            self.pool1,  # 第一个池化层
            self.conv3,  # 第三个卷积层
            self.relu3,  # 第三个ReLU激活
            self.conv4,  # 第四个卷积层
            self.relu4,  # 第四个ReLU激活
            self.pool2,  # 第二个池化层
            self.conv5,  # 第五个卷积层
            self.relu5,  # 第五个ReLU激活
            self.pool3   # 第三个池化层
        )
        
        # 使用一个虚拟输入样本计算特征图大小
        # 这是一种动态确定全连接层输入维度的方法
        dummy_input = torch.zeros(1, 3, 224, 224)  # 创建一个1x3x224x224的全零张量
        with torch.no_grad():  # 不计算梯度
            features_output = self.features(dummy_input)  # 通过特征提取网络
        
        # 计算展平后的特征维度 (批次大小 x 特征数)
        self.flat_features = features_output.size(1) * features_output.size(2) * features_output.size(3)
        print(f"CNN特征维度: {self.flat_features} (通道数: {features_output.size(1)}, 高: {features_output.size(2)}, 宽: {features_output.size(3)})")
        
        # 全连接层部分 - 第一个全连接层 + ReLU
        # 输入: 展平的特征 -> 输出: 4096维向量
        self.fc6 = nn.Linear(self.flat_features, 4096)
        self.relu6 = nn.ReLU(inplace=True)
        
        # 全连接层部分 - 第二个全连接层 + ReLU
        # 输入: 4096维向量 -> 输出: 1000维向量
        self.fc7 = nn.Linear(4096, 1000)
        self.relu7 = nn.ReLU(inplace=True)
        
        # 全连接层部分 - 输出层
        # 输入: 1000维向量 -> 输出: 类别数维向量
        self.fc8 = nn.Linear(1000, num_classes)
        
        # Dropout层 - 随机丢弃50%的神经元，防止过拟合
        # 在训练时随机将一些神经元输出置为0，测试时不起作用
        self.dropout = nn.Dropout(p=0.5)  # p=0.5表示丢弃概率为50%

    def forward(self, x):
        """
        前向传播
        
        参数:
            x (torch.Tensor): 输入张量
            
        返回:
            torch.Tensor: 模型输出
        """
        # 特征提取
        x = self.features(x)
        
        # 展平，为全连接层做准备
        x = x.view(x.size(0), -1)
        
        # fc6 -> relu6 -> dropout
        x = self.dropout(self.relu6(self.fc6(x)))
        
        # fc7 -> relu7 -> dropout
        x = self.dropout(self.relu7(self.fc7(x)))
        
        # fc8
        x = self.fc8(x)
        
        return x 